"""ZenML Pipeline for CrewAI Travel Assistant.

This pipeline encapsulates the CrewAI multi-agent system in a ZenML pipeline,
demonstrating how to integrate the CrewAI framework with ZenML for orchestration
and artifact management.
"""

import os
from typing import Annotated, Any, Dict

from crewai import Agent, Crew, Task
from crewai.tools import tool

from zenml import pipeline, step
from zenml.config import DockerSettings, PythonPackageInstaller

docker_settings = DockerSettings(
    python_package_installer=PythonPackageInstaller.UV,
    requirements="requirements.txt",  # relative to the pipeline directory
    environment={
        "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY"),
        # Set home directory to a writable location for CrewAI storage
        "HOME": "/tmp",  # nosec B108 - Docker env var, not insecure file operation
        # Alternative: override the specific appdirs behavior
        "XDG_DATA_HOME": "/tmp/.local/share",  # nosec B108 - Docker env var, not insecure file operation
    },
)


@tool("Weather Checker Tool")
def get_weather(city: str) -> str:
    """Get weather for a given city."""
    return f"Current weather in {city}: Sunny, 22°C (72°F), light breeze"


# Weather Specialist Agent
weather_checker = Agent(
    role="Weather Specialist",
    goal="Check weather conditions for a given city: {city}",
    backstory="You are a meteorology expert who provides accurate weather updates. When you get weather information, you immediately report it clearly.",
    tools=[get_weather],
    verbose=True,
    max_iter=5,
)


# Travel Advisor Agent
travel_advisor = Agent(
    role="Travel Advisor",
    goal="Give practical travel advice based on weather conditions for {city}",
    backstory="You are an experienced travel advisor who helps people prepare for their trips.",
    verbose=True,
    max_iter=5,
)


# Task 1: Check weather with parameterized city
check_weather_task = Task(
    description="Check the current weather in {city} and provide a weather report.",
    expected_output="A clear statement of the current weather conditions in {city} including temperature and conditions",
    agent=weather_checker,
)


# Task 2: Packing advice based on weather
packing_advice_task = Task(
    description="Based on the weather report for {city}, provide 3-5 specific items to pack for the trip.",
    expected_output="A list of 3-5 specific items to pack based on the weather conditions",
    agent=travel_advisor,
    context=[check_weather_task],
)


# PanAgent will discover this 'crew' variable
crew = Crew(
    agents=[weather_checker, travel_advisor],
    tasks=[check_weather_task, packing_advice_task],
    verbose=True,
)


@step
def run_crewai_agents(city: str) -> Annotated[Dict[str, Any], "crew_results"]:
    """Execute the CrewAI crew and return results."""
    try:
        # Execute the crew with the city parameter
        result = crew.kickoff(inputs={"city": city})

        # Convert result to dict for ZenML artifact storage
        return {"city": city, "result": str(result), "status": "success"}
    except Exception as e:
        return {
            "city": city,
            "result": f"Crew execution error: {str(e)}",
            "status": "error",
        }


@step
def format_travel_results(
    crew_data: Dict[str, Any],
) -> Annotated[str, "formatted_results"]:
    """Format the CrewAI results into a readable summary."""
    city = crew_data["city"]
    result = crew_data["result"]
    status = crew_data["status"]

    if status == "error":
        formatted = f"""❌ TRAVEL PLANNING ERROR FOR {city.upper()}
{"=" * 50}

{result}
"""
    else:
        formatted = f"""✈️ TRAVEL PLANNING FOR {city.upper()}
{"=" * 50}

{result}

🤖 Generated by CrewAI Multi-Agent System
"""

    return formatted.strip()


@pipeline(settings={"docker": docker_settings}, enable_cache=False)
def agent_pipeline(city: str = "Berlin") -> str:
    """ZenML pipeline that orchestrates the CrewAI travel planning system.

    Returns:
        Formatted travel planning results
    """
    # Run the CrewAI agents
    crew_results = run_crewai_agents(city=city)

    # Format the results
    summary = format_travel_results(crew_results)

    return summary


if __name__ == "__main__":
    print("🚀 Running CrewAI travel pipeline...")
    run_result = agent_pipeline()
    print("Pipeline completed successfully!")
    print("Check the ZenML dashboard for detailed results and artifacts.")
